#####
# Replication for: "Can Political Speech Foster Tolerance of Immigrants?" by Schleiter, Tavits, and Ward.
# Figure S.5
#####

library(here)
library(data.table)
library(ggplot2)

# source the custom functions
source("functions.R")

# load the pooled data if not already in workspace
if(!exists("pooled")){
  pooled <- fread(here("data", "pooled.csv"))
}

# coerce the heterogeneity measure to a factor if not already
if(class(pooled$pid2_het) != "factor"){
  pooled[ , pid2_het := factor(pid2_het, levels = c("Independent", "Democrat", "Republican"))]
}


# Helper functions --------------------------------------------------------

# These three functions do the starch of the analysis. The first fits all 9 models, the second extracts the 9 marginal effects per model. The third does the battery of p-values
allCovPIDfits <- 
  function(y, samp, data = pooled, chatty = T){
    if(samp != "Replication"){
      form <- as.formula(paste0(y, "~ CH + NM + CS + CH*pid2_het + NM*pid2_het + CS*pid2_het + female + age + I(age^2) + race2 + edu2 + news2 + region2 + factor(wave)"))
    } else {
      form <- as.formula(paste0(y, "~ CH + NM + CS + CH*pid2_het + NM*pid2_het + CS*pid2_het + female + age + I(age^2) + race2 + edu2 + news2 + region2"))
    }
    
    sample_use <- c("Main" = "replication == 0", "Replication" = "replication == 1", "Pooled" = "")[samp]
    
    if(samp != "Pooled"){
      mod <- lm(formula = form, data = data[eval(parse(text = sample_use))])
    } else {
      mod <- lm(formula = form, data = data)
    }
    
    if(chatty) print(summary(mod))
    return(mod)
  }

PID_ME_extract <- 
  function(mod, level = .95){
    out <- expand.grid("treatment" = c("CH", "NM", "CS"), "PID" = c("I", "D", "R"), stringsAsFactors = F); setDT(out)
    
    out[PID == "I" , c("est", "lower", "upper") := .(
      coef(mod)[c("CH", "NM", "CS")],
      confint(mod, level = level)[c("CH", "NM", "CS"), 1],
      confint(mod, level = level)[c("CH", "NM", "CS"), 2]
    )]
    
    for(n in 4:9){
      vars <- c(out$treatment[n], paste0(out$treatment[n], ":pid2_het", ifelse(out$PID[n] == "D", "Democrat", "Republican")))
      est <- sum(coef(mod)[vars])
      se <- sqrt(vcov(mod)[vars[1], vars[1]] + vcov(mod)[vars[2], vars[2]] + 2*vcov(mod)[vars[1], vars[2]])
      set(out, i = n, j = "est", value = est)
      set(out, i = n, j = "lower", value = est - 1.96*se)
      set(out, i = n, j = "upper", value = est + 1.96*se)
    }
    
    return(out)
  }

# Fitting -----------------------------------------------------------------

pid_fits <- expand.grid(
  "y" = c("immPCA", "imm_neighbors2", "imm_increase2"),
  "samp" = c("Main", "Replication", "Pooled")
  )
setDT(pid_fits)

pid_fits[ , fit := Map(allCovPIDfits, y, samp, chatty = F)]
pid_me <- pid_fits[ , PID_ME_extract(mod = fit[[1]]), by = .(y, samp)]


# Plotting prep -----------------------------------------------------------

pid_me[ , yval := c("CH" = 1L, "CS" = 2L, "NM" = 3L)[treatment]]
pid_me[ , PID := factor(PID, labels = c("Democrat", "Independent", "Republican"))]
pid_me[ , y2 := factor(y, labels = c("Imm. Index (SD = 1)", "Imm. Neighbors (1-7)", "Increase Imm. (1-7)"))]
pid_me[ , samp2 := factor(samp, labels = c("Main Sample", "Replication Sample", "Pooled Sample"))]


# Plotting and saving -----------------------------------------------------

PID_plot <- ggplot(pid_me, aes(x = est, xmax = upper, xmin = lower, y = yval, color = PID, shape = PID)) +
  geom_vline(xintercept = 0, linetype = "dashed") +
  ggstance::geom_pointrangeh(position = ggstance::position_dodge2v(.7, reverse = T)) +
  facet_grid(y2 ~ samp2) +
  scale_y_continuous(breaks = c(1,2,3), labels = c("Common\nHumanity", "Countering\nStereotypes", "Norms"), limits = c(.5, 3.5), expand = expansion(0,0)) +
  ylab(NULL) +
  xlab("Estimated Treatment Effect") +
  scale_color_grey(name = "Party ID", end = 0.7) +
  scale_shape(name = "Party ID") +
  theme_bw() +
  theme(legend.position = "bottom",
        axis.title.x = element_text(size = 10),
        axis.text = element_text(size = 10),
        strip.text = element_text(size = 10),
        legend.text = element_text(size = 10),
        legend.title = element_text(size = 10))

SI_width = 6.50127  
ggsave(here("SI", "figure_s5.pdf"), PID_plot, width = SI_width, height = SI_width*1.25)

